import cn.wxw.ai.MilvusVectorStoreV2;
import io.milvus.pool.MilvusClientV2Pool;
import io.milvus.pool.PoolConfig;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.BatchingStrategy;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.time.Duration;

/**
 * @description: milvus 配置类
 * @author: wxw
 * @date: 2025/9/8
 */
@Configuration
public class MilvusConfig {

    @Value("${milvus.url}")
    private String url;

    @Value("${milvus.token}")
    private String token;

    @Value("${milvus.config.maxIdlePerKey}")
    private int maxIdlePerKey;

    @Value("${milvus.config.maxTotalPerKey}")
    private int maxTotalPerKey;

    @Value("${milvus.config.maxTotal}")
    private int maxTotal;

    @Value("${milvus.config.maxBlockWaitDuration}")
    private long maxWaitDuration;

    @Value("${milvus.config.minEvictableIdleDuration}")
    private long minEvictableIdleDuration;

    @Bean
    public MilvusClientV2 milvusClientV2() throws ClassNotFoundException, NoSuchMethodException {
        ConnectConfig connectConfig = ConnectConfig.builder()
                .uri(url)
                .token(token)
                .build();

        PoolConfig poolConfig = PoolConfig.builder()
                .maxIdlePerKey(maxIdlePerKey) // max idle clients per key
                .maxTotalPerKey(maxTotalPerKey) // max total(idle + active) clients per key
                .maxTotal(maxTotal) // max total clients for all keys
                .maxBlockWaitDuration(Duration.ofSeconds(maxWaitDuration)) // getClient() will wait 5 seconds if no idle client available
                .minEvictableIdleDuration(Duration.ofSeconds(minEvictableIdleDuration)) // if number of idle clients is larger than maxIdlePerKey, redundant idle clients will be evicted after 10 seconds
                .build();
        MilvusClientV2Pool pool = new MilvusClientV2Pool(poolConfig, connectConfig);
        // TODO: client_name 的设置暂不清楚, 目前是设置成了节点名称
        return pool.getClient("proxy13");
    }

    // 定义 MilvusVectorStoreV2 Bean
    @Bean
    public MilvusVectorStoreV2 milvusVectorStoreV2(
            MilvusClientV2 milvusClientV2,
            OpenAiEmbeddingModel openAiEmbeddingModel,
            BatchingStrategy batchingStrategy) {

        return new MilvusVectorStoreV2(
                milvusClientV2,
                openAiEmbeddingModel,
                batchingStrategy
        );
    }

}

